#!/usr/bin/env python
from __future__ import print_function
import os
import subprocess
from maputility2 import maptools
from glat2 import mdltk

def gtmodel(srcMaps, expCube, binnedExpMap, srcModel, irfs, evtype, outfile, edisp_bins=-2):
    cmdline = ['gtmodel', 'irfs='+irfs, 'outfile='+outfile,
               'srcmaps='+srcMaps, 'expcube='+expCube, 'bexpmap='+binnedExpMap, 'srcmdl='+srcModel,
               'outtype=ccube', 'evtype=%d'%evtype, 'edisp_bins=%d'%edisp_bins]
    print(' '.join(cmdline))
    subprocess.check_call(cmdline)

def make_binned_mcube(srcmaps, ltcube, expcube, srcmdl_first, irfs, evtype, outfile, edisp_bins=-2):
    """
    srcmdl_first is the first xmlfile, the program will automatically find the subsequent ones
    according to the naming scheme of the file (such as model_pass2_band00_wid02.xml)
    [PREFIX]_band[STARTBIN]_wid[BINWIDTH].xml
    or
    [PREFIX]_BAND[STARTBIN].xml (assume binwidth=1)
    """
    mcube = maptools.mapLike(srcmaps, fill_value=np.nan) # use nan to notify the data outside the E range
    ebins = mdltk.EnergyBins(srcmaps)

    path, srcmdl = os.path.split(srcmdl_first)
    if not path: path = './'

    files_raw = os.listdir(path)
    files = []
    for f in files_raw:
        if f.endswith('.xml'):
            files.append(f)

    while srcmdl is not None:
        fn0 = srcmdl.split('.')[0]
        if 'wid' in fn0:
            fn0, suff = fn0.split('_wid')
            widbin = int(suff)
        else:
            widbin = 1

        nextsrcmdl = None
        if 'band' in fn0:
            pref, suff = fn0.split('band')
            idxmin = int(suff)
            idxmax = idxmin + widbin

            pref_next = pref+'band%02d'%idxmax
            for f in files:
                if f.startswith(pref_next):
                    nextsrcmdl = f
        else:
            idxmin, idxmax = 0, mcube.nlayers

        srcmdl_full = os.path.join(path, srcmdl)
        gtmodel(srcMaps=srcmaps, expCube=ltcube, binnedExpMap=expcube, srcModel=srcmdl_full,
                irfs=irfs, evtype=evtype, outfile=outfile, edisp_bins=edisp_bins)

        mcube0 = maptools.mapOpen(outfile)
        mcube.mapdata[idxmin:idxmax,:,:] = mcube0.mapdata[idxmin:idxmax,:,:]

        srcmdl = nextsrcmdl

    mcube.saveMap(outfile)
    ebins.append_hdu_tofile(outfile)

def cli():
    import argparse

    parser = argparse.ArgumentParser(description='make mcube according to the srcmdl')
    parser.add_argument("srcMaps", type=str, help='Source Maps')
    parser.add_argument("expCube", type=str, help='Live Time Cube')
    parser.add_argument("binnedExpMap", type=str, help='Exposure Cube')
    parser.add_argument("srcModelFIRST", type=str, help='model_pass2_band00_wid02.xml')
    parser.add_argument("irfs", type=str, help='IRFs')
    parser.add_argument("evtype", type=int, help='event type')
    parser.add_argument("outfile", type=str, help='output mcube file')
    parser.add_argument("-edispbins", type=int, default=-2, help='the extra ebins added (default=-2)')
    args = parser.parse_args()

    make_binned_mcube(
        srcmaps=args.srcMaps, ltcube=args.expCube, expcube=args.binnedExpMap,
        srcmdl_first=args.srcModelFIRST, irfs=args.irfs, evtype=args.evtype,
        outfile=args.outfile, edisp_bins=args.edispbins)

if __name__ == '__main__': cli()
